
from pymongo import ASCENDING
from pev3.util.DataBase import  DB_CONN
from datetime import datetime, timedelta

"""
获取指定日期范围的按照正序排列的交易日列表
如果没有指定日期范围，则获取从当期日期向前365个自然日内的所有交易日

:param startDate: 开始日期
:param endDate: 结束日期
:return: 日期列表
"""
def getTradingDates(startDate=None, endDate=None):

    # 开始日期，默认今天向前的365个自然日
    now = datetime.now()

    if startDate is None:
        onYearAgo = now - timedelta(days=365)
        startDate = onYearAgo.strftime('%Y-%m-%d')

    # 结束日期默认为今天
    if endDate is None:
        endDate = now.strftime('%Y-%m-%d')

    dailyCursor = DB_CONN.daily.find(
        {'code': '000001', 'date': {'$gte': startDate, '$lte': endDate}, 'index': True},
        sort=[('date', ASCENDING)],
        projection={'date': True, '_id':False}
    )

    dates = [x['date'] for x in dailyCursor]

    return dates

"""
获取某个交易日的所有股票代码列表，如果没有指定日期，则从当前日期一直向前找，直到找到有
数据的一天，返回的即是那个交易日的股票代码列表

:param date: 日期
:return: 股票代码列表
"""
def getTradingDateStockCodes(date=None):
    datetimeObj = datetime.now()

    if date is None:
        date = datetimeObj.strftime('%Y-%m-%d')

    codes = []
    while len(codes) == 0:
        codeCursor = DB_CONN.basic.find(
            {'date': date},
            projection={'code': True, '_id': False}
        )

        codes = [x['code'] for x in codeCursor]

        datetimeObj = datetimeObj - timedelta(days=1)
        date = datetimeObj.strftime('%Y-%m-%d')

    return codes